package org.anno.chunkdb;

import com.google.common.base.Preconditions;

import it.unimi.dsi.fastutil.ints.IntArrayList;
import it.unimi.dsi.fastutil.ints.IntIterator;
import it.unimi.dsi.fastutil.ints.IntList;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Comparator;
import java.util.List;
import java.util.TreeMap;
import java.util.stream.Collectors;

/**
 * Builds {@code byte[]} for a table.
 */
public class ChunkBuilder {

  private final List<ColumnBuilder> columnBuilders;
  private int size;
  private int rowCount;

  public ChunkBuilder(List<Metadata.Column> columns) {
    columnBuilders = columns.stream().map(this::make).collect(Collectors.toList());
    // 4 = 2 version + 2 rowCount
    size = 4 + 3 * columns.size();
  }

  private ColumnBuilder make(Metadata.Column column) {
    return new ColumnBuilder();
  }

  public ChunkBuilder addRow(List<String> row) {
    Preconditions.checkArgument(row.size() == columnBuilders.size());
    for (int i = 0; i < columnBuilders.size(); i++) {
      size += columnBuilders.get(i).addCell(row.get(i));
    }
    rowCount += 1;
    return this;
  }

  public byte[] build() {
    ByteArrayOutputStream baos = new ByteArrayOutputStream(1 << 20);
    // int16 version = 1
    baos.write(0);
    baos.write(1);
    int[] columnStart = new int[columnBuilders.size()];
    int i = 0;
    for (ColumnBuilder builder : columnBuilders) {
      columnStart[i++] = baos.size();
      builder.appendTo(baos);
    }
    for (int p : columnStart) {
      baos.write(p >> 16);
      baos.write(p >> 8);
      baos.write(p);
    }
    baos.write(rowCount >> 8);
    baos.write(rowCount);
    return baos.toByteArray();
  }

  public int getSize() {
    return size;
  }

  public int getEstimatedSizeAfterAdding(String[] line) {
    // 10 = 5 value start + 2 index + 3 value length
    int maxSize = 10 + columnBuilders.size();
    for (String s : line) {
      maxSize += 2 * s.length();
    }
    return maxSize;
  }

  private static class ColumnBuilder {

    private static final Comparator<byte[]> BYTE_ARRAY_ORDER =
        (o1, o2) -> FilterQueryPlan.compare(o1, 0, o1.length, o2);

    private final ByteArrayOutputStream baos = new ByteArrayOutputStream();
    private final IntList valueOffset = new IntArrayList();
    private final TreeMap<byte[], IntList> index = new TreeMap<>(BYTE_ARRAY_ORDER);

    public int addCell(String cell) {
      return addBytes(cell.getBytes(StandardCharsets.UTF_8));
    }

    protected int addBytes(byte[] bytes) {
      IntList list = index.get(bytes);
      if (list == null) {
        list = new IntArrayList();
        index.put(bytes, list);
      }
      list.add(valueOffset.size());
      int initialSize = baos.size();
      valueOffset.add(initialSize);
      writeVarInt(bytes.length, baos);
      baos.write(bytes, 0, bytes.length);
      return baos.size() - initialSize + 2 + ((valueOffset.size() & 1) == 1 ? 5 : 0);
    }

    void appendTo(ByteArrayOutputStream writer) {
      int start = writer.size();
      try {
        baos.writeTo(writer);
      } catch (IOException e) {
        throw new AssertionError(e);
      }
      for (int i = 0; i < valueOffset.size(); i += 2) {
        int valueStart = start + valueOffset.getInt(i);
        writer.write(valueStart >> 12);
        writer.write(valueStart >> 4);
        int split = (valueStart & 15) << 4;
        valueStart = i + 1 < valueOffset.size() ? valueOffset.getInt(i + 1) + start : 0;
        split |= ((valueStart >> 16) & 15);
        writer.write(split);
        writer.write(valueStart >> 8);
        writer.write(valueStart);
      }
      for (IntList intList : index.values()) {
        for (IntIterator iterator = intList.iterator(); iterator.hasNext(); ) {
          int row = iterator.nextInt();
          writer.write(row >> 8);
          writer.write(row);
        }
      }
    }
  }

  public static void writeVarInt(int length, ByteArrayOutputStream out) {
    for (int i = 28; i > 0; i -= 7) {
      if (length >= (1 << i)) {
        out.write(128 | (127 & (length >> (i - 7))));
      }
    }
    out.write(length & 127);
  }
}
